import math
import gymnasium as gym
from gymnasium import spaces
import pybullet as p
import pybullet_data
import numpy as np

class BipedalWheeledRobotEnv(gym.Env):
    def __init__(self, render_mode=None):
        super(BipedalWheeledRobotEnv, self).__init__()

        # Store the render mode if provided
        self.render_mode = render_mode

        # Connect to the PyBullet simulator
        self.client = p.connect(p.GUI)
        p.setAdditionalSearchPath(pybullet_data.getDataPath())
        p.setGravity(0, 0, -9.81)

        # Load the plane and robot URDF
        self.plane = p.loadURDF("plane.urdf")
        self.robot = p.loadURDF("bipedal_wheeled_robot.urdf", [0, 0, 1], useFixedBase=False)
        
        # Define a mapping from joint names to joint IDs (direct mapping from joint names to corresponding IDs)
        self.joint_name_to_id = {
            "body_to_left_thigh": 0,
            "left_thigh_to_left_thigh2": 1,
            "body_to_right_thigh": 2,
            "right_thigh_to_right_thigh2": 3,
            "left_thigh2_to_left_wheel": 4,
            "right_thigh2_to_right_wheel": 5
        }

        # State space: pitch angle (rad), wheel contact position (m), body angular velocity (rad/s), wheel contact velocity (m/s)
        self.observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=(3,), dtype=np.float32)

        # Action space: control left and right wheel speeds, range [-1, 1] m/s
        self.action_space = spaces.Box(low=-1, high=1, shape=(1,), dtype=np.float32)

        # Initialization
        self.target_velocity = 0.5
        self.max_steps = 1000
        self.current_step = 0
        self.angular_weight = 1
        self.velocity_weight = 0.4



        # Servo-related initialization (removed wheel servo-related parts)
        self.leg_return_period = 1.0  # Time constant for the hip and knee servo motors to return to the neutral position
        self.wheel_radius = 0.2  # Wheel radius, in meters

        # Store servo actions
        self.__leg_servo_action = {
            joint: {
                "position": 0.0,
                "velocity": 0.0,
                "maximum_torque": 100,  # Assumed maximum torque
            }
            for joint in ["body_to_left_thigh", "left_thigh_to_left_thigh2", "body_to_right_thigh", "right_thigh_to_right_thigh2"]  # Assume names for hip and knee joints
        }

    def reset(self, seed=None, options=None):
        """Reset the environment and add a random seed"""
        self.current_step = 0

        super().reset(seed=seed)


        p.resetSimulation()
        p.setGravity(0, 0, -9.81)
        self.plane = p.loadURDF("plane.urdf")
        self.robot = p.loadURDF("bipedal_wheeled_robot.urdf", [0, 0, 1], useFixedBase=False)

        # Return the initial observation
        obs = self.get_observation()

        # Return (observation, empty info dictionary)
        return obs, {}

    def set_initial_wheel_velocity(self, velocity):
        """Set the initial speed of the wheels"""
        p.setJointMotorControl2(self.robot, 4, p.VELOCITY_CONTROL, targetVelocity=velocity[0])
        p.setJointMotorControl2(self.robot, 5, p.VELOCITY_CONTROL, targetVelocity=velocity[0])

    def get_upper_leg_servo_action(self):
        """Get the servo motor actions for the hip and knee joints"""
        for joint in self.__leg_servo_action.keys():
            # Use pybullet to get the current state of the joint (position, velocity, torque, etc.)
            joint_id = self.joint_name_to_id[joint]
            joint_state = p.getJointState(self.robot, joint_id)
            current_position = joint_state[0]  # Current position of the joint
        
            # Update the servo motor position, gradually return to the neutral position (0)
            new_position = self.low_pass_filter(
                prev_output=current_position,  # Use the current joint position as the initial output
                cutoff_period=self.leg_return_period,
                new_input=0.0,  # Target is to return to the neutral position 0
                dt=1/240  # Default timestep in PyBullet
            )
        
            # Update the servo motor action
            self.__leg_servo_action[joint]["position"] = new_position
        return self.__leg_servo_action

    def low_pass_filter(self, prev_output, cutoff_period, new_input, dt):
        """Simple low-pass filter"""
        alpha = dt / (cutoff_period + dt)

        return alpha * new_input + (1.0 - alpha) * prev_output

    def step(self, action):
        """Execute action, return (observation, reward, done, truncated, info)"""
        self.current_step += 1

        # Set left and right wheel speeds based on action (directly controlled by reinforcement learning)
        wheel_velocity = action[0]/self.wheel_radius  # Wheel speed
        #print(wheel_velocity)
        
        # Get servo motor actions
        servo_dict = self.get_upper_leg_servo_action()

        # Control joint actions (assuming hip and knee joint IDs are 0 - 3, wheel joint IDs are 4 and 5)
        p.setJointMotorControl2(self.robot, 0, p.POSITION_CONTROL, targetPosition=servo_dict["body_to_left_thigh"]["position"])
        p.setJointMotorControl2(self.robot, 1, p.POSITION_CONTROL, targetPosition=servo_dict["left_thigh_to_left_thigh2"]["position"])
        p.setJointMotorControl2(self.robot, 2, p.POSITION_CONTROL, targetPosition=servo_dict["body_to_right_thigh"]["position"])
        p.setJointMotorControl2(self.robot, 3, p.POSITION_CONTROL, targetPosition=servo_dict["right_thigh_to_right_thigh2"]["position"])

        # Directly set wheel speed control
        p.setJointMotorControl2(self.robot, 4, p.VELOCITY_CONTROL, targetVelocity=wheel_velocity)
        p.setJointMotorControl2(self.robot, 5, p.VELOCITY_CONTROL, targetVelocity=wheel_velocity)

        # Step the simulation
        p.stepSimulation()

        # Get new observation values
        obs = self.get_observation()

        # Calculate reward function
        reward = self.calculate_reward(obs)

        # Check if the task is done (e.g., if the robot has fallen)
        terminated = self.check_if_robot_fell(obs)
        truncated = self.current_step >= self.max_steps  # Check if the maximum timestep has been reached

        # Return five values: observation, reward, done, truncated, info
        info = {}
        return obs, reward, terminated, truncated, info

    def get_observation(self):
        """Get the robot's state as an observation"""
        pos, orn = p.getBasePositionAndOrientation(self.robot)
        vel, ang_vel = p.getBaseVelocity(self.robot)

        euler = p.getEulerFromQuaternion(orn)
        pitch_angle = euler[1]  # Pitch angle
        wheel_contact_velocity = vel[0]  # Velocity in the X direction

        state = np.array([
            pitch_angle,
            ang_vel[1],  # Angular velocity around the Y axis
            wheel_contact_velocity
        ])
        return state

    def check_if_robot_fell(self, obs):
        pitch_angle = obs[0]
        # Check if the pitch angle exceeds a certain threshold (assume robot falls if pitch angle exceeds 45 degrees)
        if abs(pitch_angle) > np.pi / 4:  # 45 degrees threshold
            return True
        # Not fallen
        return False

    def calculate_reward(self, obs):
        """Calculate the reward, including penalty for target velocity deviation"""
        pitch_angle, body_angular_velocity, wheel_contact_velocity = obs
         
        top_height = 1 

        top_velocity = wheel_contact_velocity + top_height * body_angular_velocity * np.cos(pitch_angle)

        # Use squared penalty for a smoother gradient
        velocity_penalty = -((top_velocity - self.target_velocity) ** 2)  

        # Calculate angular reward, scaling factors can be adjusted
        angular_reward = np.exp(-np.abs(body_angular_velocity)) + np.exp(-10*np.abs(pitch_angle))


        # Combine rewards with weights
        reward = (self.angular_weight * angular_reward + 
              self.velocity_weight * velocity_penalty )
        
        return reward

    def render(self):
        if self.render_mode is None:
            print("You tried to call render() but no `render_mode` was passed to the env constructor.")
            return
        
    def close(self):
        p.disconnect(self.client)
